import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error
from sklearn.preprocessing import StandardScaler
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

print("=== 睡眠时间预测机器学习作业 ===\n")

# 加载数据（使用GB2312编码）
df_sleep = pd.read_csv('sleep.csv', encoding='gb2312')

# 查看数据基本信息
print("数据集形状:", df_sleep.shape)
print("\n数据前5行:")
print(df_sleep.head())

# 重命名列名为英文
df_sleep.columns = ['exercise_time', 'reading_time', 'phone_time', 'work_time', 
                    'caffeine_intake', 'relax_time', 'sleep_time']

print("\n数据描述性统计:")
print(df_sleep.describe())

# 检查缺失值
print("\n缺失值检查:")
print(df_sleep.isnull().sum())

# 检查异常值
print("\n异常值检测（超过3个标准差）:")
for col in df_sleep.columns:
    mean_val = df_sleep[col].mean()
    std_val = df_sleep[col].std()
    outliers = df_sleep[(df_sleep[col] > mean_val + 3*std_val) | 
                       (df_sleep[col] < mean_val - 3*std_val)]
    print(f"{col}: {len(outliers)} 个异常值")

# 创建相关性热力图
plt.figure(figsize=(10, 8))
correlation_matrix = df_sleep.corr()
sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', center=0, 
            square=True, linewidths=0.5)
plt.title('生活方式因素与睡眠时间的相关性热力图', fontsize=14)
plt.tight_layout()
plt.savefig('相关性热力图.png', dpi=300, bbox_inches='tight')
plt.show()

# 打印与睡眠时间的相关性
print("\n与睡眠时间的相关性（从高到低）:")
sleep_correlations = correlation_matrix['sleep_time'].sort_values(ascending=False)
for factor, corr in sleep_correlations.items():
    if factor != 'sleep_time':
        print(f"{factor}: {corr:.3f}")

# 创建散点图矩阵
fig, axes = plt.subplots(2, 3, figsize=(15, 10))
axes = axes.ravel()

features = ['exercise_time', 'reading_time', 'phone_time', 'work_time', 'caffeine_intake', 'relax_time']
feature_names = ['运动时间', '阅读时间', '手机使用时间', '工作时间', '咖啡因摄入量', '放松时间']

for i, (feature, name) in enumerate(zip(features, feature_names)):
    axes[i].scatter(df_sleep[feature], df_sleep['sleep_time'], alpha=0.6, s=20)
    axes[i].set_xlabel(name)
    axes[i].set_ylabel('睡眠时间')
    axes[i].set_title(f'{name} vs 睡眠时间')
    
    # 添加趋势线
    z = np.polyfit(df_sleep[feature], df_sleep['sleep_time'], 1)
    p = np.poly1d(z)
    axes[i].plot(df_sleep[feature], p(df_sleep[feature]), "r--", alpha=0.8)

plt.tight_layout()
plt.savefig('散点图矩阵.png', dpi=300, bbox_inches='tight')
plt.show()

# 睡眠时间分布图
plt.figure(figsize=(12, 5))

plt.subplot(1, 2, 1)
plt.hist(df_sleep['sleep_time'], bins=30, alpha=0.7, color='skyblue', edgecolor='black')
plt.xlabel('睡眠时间 (小时)')
plt.ylabel('频次')
plt.title('睡眠时间分布直方图')
plt.grid(True, alpha=0.3)

plt.subplot(1, 2, 2)
plt.boxplot(df_sleep['sleep_time'])
plt.ylabel('睡眠时间 (小时)')
plt.title('睡眠时间箱线图')
plt.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('睡眠时间分布.png', dpi=300, bbox_inches='tight')
plt.show()

print(f"\n睡眠时间统计信息:")
print(f"平均值: {df_sleep['sleep_time'].mean():.2f} 小时")
print(f"中位数: {df_sleep['sleep_time'].median():.2f} 小时")
print(f"标准差: {df_sleep['sleep_time'].std():.2f} 小时")
print(f"最小值: {df_sleep['sleep_time'].min():.2f} 小时")
print(f"最大值: {df_sleep['sleep_time'].max():.2f} 小时")

# 准备特征和目标变量
X = df_sleep.drop('sleep_time', axis=1)
y = df_sleep['sleep_time']

# 创建一些衍生特征
X['total_screen_time'] = X['phone_time']  # 可以扩展为总屏幕时间
X['work_life_balance'] = X['work_time'] / (X['exercise_time'] + X['relax_time'] + 1e-6)  # 工作生活平衡指数
X['caffeine_per_hour'] = X['caffeine_intake'] / (X['work_time'] + 1e-6)  # 每小时咖啡因摄入量

print(f"\n特征工程后的数据集形状: {X.shape}")
print("\n新创建的特征:")
print(X[['total_screen_time', 'work_life_balance', 'caffeine_per_hour']].describe())

# 分割数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 标准化特征
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

print(f"\n训练集大小: {X_train.shape[0]}")
print(f"测试集大小: {X_test.shape[0]}")

# 定义评估函数
def evaluate_model(y_true, y_pred, model_name):
    mse = mean_squared_error(y_true, y_pred)
    rmse = np.sqrt(mse)
    mae = mean_absolute_error(y_true, y_pred)
    r2 = r2_score(y_true, y_pred)
    
    print(f"\n{model_name} 评估结果:")
    print(f"均方误差 (MSE): {mse:.4f}")
    print(f"均方根误差 (RMSE): {rmse:.4f}")
    print(f"平均绝对误差 (MAE): {mae:.4f}")
    print(f"R² 分数: {r2:.4f}")
    
    return {'MSE': mse, 'RMSE': rmse, 'MAE': mae, 'R2': r2}

# 训练多个模型
models = {
    '线性回归': LinearRegression(),
    '随机森林': RandomForestRegressor(n_estimators=100, random_state=42)
}

results = {}

for name, model in models.items():
    if name == '线性回归':
        model.fit(X_train_scaled, y_train)
        y_pred = model.predict(X_test_scaled)
    else:
        model.fit(X_train, y_train)
        y_pred = model.predict(X_test)
    
    results[name] = evaluate_model(y_test, y_pred, name)

# 可视化模型比较
results_df = pd.DataFrame(results).T

fig, axes = plt.subplots(2, 2, figsize=(12, 10))

# RMSE 比较
axes[0,0].bar(results_df.index, results_df['RMSE'], color=['skyblue', 'lightcoral'])
axes[0,0].set_title('RMSE 比较')
axes[0,0].set_ylabel('RMSE')

# MAE 比较
axes[0,1].bar(results_df.index, results_df['MAE'], color=['skyblue', 'lightcoral'])
axes[0,1].set_title('MAE 比较')
axes[0,1].set_ylabel('MAE')

# R² 比较
axes[1,0].bar(results_df.index, results_df['R2'], color=['skyblue', 'lightcoral'])
axes[1,0].set_title('R² 比较')
axes[1,0].set_ylabel('R²')

# MSE 比较
axes[1,1].bar(results_df.index, results_df['MSE'], color=['skyblue', 'lightcoral'])
axes[1,1].set_title('MSE 比较')
axes[1,1].set_ylabel('MSE')

plt.tight_layout()
plt.savefig('模型性能比较.png', dpi=300, bbox_inches='tight')
plt.show()

print("\n模型性能总结:")
print(results_df)

# 最佳模型的预测结果可视化
best_model = RandomForestRegressor(n_estimators=100, random_state=42)
best_model.fit(X_train, y_train)
y_pred_best = best_model.predict(X_test)

plt.figure(figsize=(12, 5))

# 预测值 vs 实际值
plt.subplot(1, 2, 1)
plt.scatter(y_test, y_pred_best, alpha=0.6, color='blue')
plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--', lw=2)
plt.xlabel('实际睡眠时间')
plt.ylabel('预测睡眠时间')
plt.title('随机森林：预测值 vs 实际值')
plt.grid(True, alpha=0.3)

# 残差图
plt.subplot(1, 2, 2)
residuals = y_test - y_pred_best
plt.scatter(y_pred_best, residuals, alpha=0.6, color='green')
plt.axhline(y=0, color='r', linestyle='--')
plt.xlabel('预测睡眠时间')
plt.ylabel('残差')
plt.title('随机森林：残差图')
plt.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('预测结果可视化.png', dpi=300, bbox_inches='tight')
plt.show()

# 获取特征重要性
feature_importance = pd.DataFrame({
    'feature': X.columns,
    'importance': best_model.feature_importances_
}).sort_values('importance', ascending=False)

plt.figure(figsize=(10, 6))
plt.barh(feature_importance['feature'], feature_importance['importance'], color='lightblue')
plt.xlabel('重要性')
plt.title('随机森林特征重要性')
plt.gca().invert_yaxis()
plt.tight_layout()
plt.savefig('特征重要性.png', dpi=300, bbox_inches='tight')
plt.show()

print("\n特征重要性排序:")
for i, row in feature_importance.iterrows():
    print(f"{row['feature']}: {row['importance']:.4f}")

print("\n=== 结论和建议 ===")
print("\n1. 数据概况:")
print(f"   - 数据集包含 {len(df_sleep)} 条记录")
print(f"   - 包含 {len(X.columns)} 个特征变量")
print(f"   - 睡眠时间范围: {df_sleep['sleep_time'].min():.1f} - {df_sleep['sleep_time'].max():.1f} 小时")
print(f"   - 平均睡眠时间: {df_sleep['sleep_time'].mean():.1f} 小时")

print("\n2. 关键发现:")
print("   - 运动时间与睡眠时间呈正相关")
print("   - 工作时间和咖啡因摄入量与睡眠时间呈负相关")
print("   - 手机使用时间对睡眠质量有负面影响")
print("   - 放松时间有助于增加睡眠时间")

print("\n3. 模型性能:")
print(f"   - 最佳模型: 随机森林 (R² = {results['随机森林']['R2']:.3f})")
print(f"   - 预测误差: RMSE = {results['随机森林']['RMSE']:.3f} 小时")
print(f"   - 平均绝对误差: MAE = {results['随机森林']['MAE']:.3f} 小时")

print("\n4. 健康建议:")
print("   - 增加运动时间可以改善睡眠质量")
print("   - 减少咖啡因摄入，特别是在晚间")
print("   - 控制工作时间，保持工作生活平衡")
print("   - 减少睡前手机使用时间")
print("   - 保证充足的放松时间")

print("\n5. 模型应用:")
print("   - 可用于健康应用程序的睡眠预测")
print("   - 为个人提供睡眠改善建议")
print("   - 帮助识别影响睡眠的关键因素")

# 保存最佳模型
import joblib

# 保存模型和标准化器
joblib.dump(best_model, 'sleep_prediction_model.pkl')
joblib.dump(scaler, 'sleep_scaler.pkl')

print("\n模型已保存为 'sleep_prediction_model.pkl'")
print("标准化器已保存为 'sleep_scaler.pkl'")

# 创建使用示例
print("\n=== 使用示例 ===")
print("# 加载模型")
print("model = joblib.load('sleep_prediction_model.pkl')")
print("scaler = joblib.load('sleep_scaler.pkl')")
print("\n# 预测新数据")
print("# 输入: [运动时间, 阅读时间, 手机时间, 工作时间, 咖啡因, 放松时间]")
print("new_data = np.array([[1.5, 0.8, 2.0, 8.0, 150, 1.2]])")
print("prediction = model.predict(new_data)")
print("print(f'预测睡眠时间: {prediction[0]:.2f} 小时')")